﻿using System;
using System.Collections.Generic;
using System.Linq;
using Allegro.TaskExecution;

namespace Allegro.Mathlib
{
    public class LatinHypercubeSampler : ISampler<double[]>
    {
        private Random _random;
        private int _nParams;
        private int _nSamples;
        private double[][] _parameterSamples;
        private int[] _currentSelection;

        bool[,] _selections;
        private int _nSelections;
        private IProbabilityDistribution[] _distributions;

        /// <summary>
        /// Create a sampler with nParameters parameters
        /// </summary>
        /// <param name="nParameters"></param>
        /// <param name="distributions"></param>
        public LatinHypercubeSampler(int nParameters, int nSamples, IProbabilityDistribution[] distributions)
        {
            if (distributions.Length != nParameters)
            {
                string msg = "Dimension of distribution array different from number of parameters given";
                Log.Error(msg, null,
                          LogEventTypes.OtherEvent);
                throw new ArgumentException(msg);
            }
            _random = new Random();
            _nParams = nParameters;
            _nSamples = nSamples;
            _distributions = distributions;

            _selections = new bool[_nParams, _nSamples];
            _currentSelection = new int[_nParams];

            _Reinit();
        }

        /// <summary>
        /// Create a sampler with nParameters parameters, each with nSamples samples
        /// </summary>
        /// <param name="nParameters">Number of parameters</param>
        /// <param name="nSamples">Number of samples per parameter</param>
        /// <param name="samples">samples[i] is an array containing nSamples for parameter i</param>
        public LatinHypercubeSampler(int nParameters, int nSamples, double[][] samples)
        {
            if (samples.Length != nParameters)
            {
                string msg = "Dimension of samples array different from number of parameters given";
                Log.Error(msg, null,
                          LogEventTypes.OtherEvent);
                throw new ArgumentException(msg);
            }
            _random = new Random();

            _parameterSamples = samples;
            _nParams = nParameters;
            _nSamples = nSamples;

            _selections = new bool[_nParams, _nSamples];
            _currentSelection = new int[_nParams];

            _Reinit();

        }

        /// <summary>
        /// Current selections of the hypercube. HyperSelections[p,s] = true if sample s is selected for parameter p
        /// </summary>
        public bool[,] HypercubeSelections
        {
            get { return _selections; }
        }
        /// <summary>
        /// CurrentSelection[p] is the sample index for parameter p
        /// </summary>
        public int[] CurrentSelection
        {
            get { return _currentSelection; }
        }

        void _Reinit()
        {
            _nSelections = 0;

            for (int p = 0; p < _nParams; p++)
            {
                for (int s = 0; s < _nSamples; s++)
                {
                    _selections[p, s] = false;
                }
            }
        }

        public void BeginSampling()
        {
            _Reinit();
        }

        public double[] NextSample()
        {
            int nSelectionsLeft = _nSamples - _nSelections;
            if (nSelectionsLeft <= 0)
            {
                _Reinit();
                nSelectionsLeft = _nSamples - _nSelections;
            }
            if (nSelectionsLeft > 1)
            {
                for (int p = 0; p < _nParams; p++)
                {
                    // Make a random selection for parameter p
                    int parameterSelection = _random.Next(nSelectionsLeft);
                    int sample = 0;
                    for (int s = 0; s < _nSamples; s++)
                    {
                        if (sample == parameterSelection && !_selections[p, s])
                        {
                            _currentSelection[p] = s;
                            _selections[p, s] = true;
                            //_sampleSelected[s] = true;
                            break;
                        }
                        if (!_selections[p, s])
                            sample++;
                    }
                }
            }
            else
            {
                // Only one selection remains, select it
                for (int p = 0; p < _nParams; p++)
                {
                    for (int s = 0; s < _nSamples; s++)
                    {
                        if (!_selections[p, s])
                        {
                            _currentSelection[p] = s;
                            _selections[p, s] = true;
                            //_sampleSelected[s] = true;
                            break;
                        }
                    }
                }
            }
            _nSelections++;

            double[] result = new double[_nParams];
            double dp = (1.0 / _nSamples);
            for (int p = 0; p < _nParams; p++)
            {
                int s = _currentSelection[p];
                if (_distributions != null)
                {
                    // Compute the probability interval to sample from
                    double pMin = dp * s;
                    double pMax = dp * (s + 1);
                    double u = _random.NextDouble();
                    double sample = pMin + u * (pMax - pMin);
                    result[p] = _distributions[p].Inverse(sample);
                }
                else if (_parameterSamples != null)
                    result[p] = _parameterSamples[p][s];
                else
                {
                    throw new ApplicationException("LatinHypercubeSampler not properly initialized. No data or distribution to sample from");
                }
            }
            return result;
        }

        public double[][] NextSamples(int n)
        {
            double[][] samples = new double[n][];
            for (int s = 0; s < n; s++)
            {
                samples[s] = NextSample();
            }
            return samples;
        }

        public bool More
        {
            get { return true; }
        }
    }
}
